#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/3/29 21:45
# @Author  : Rem~
# @File    : 3_3.py
# @function: 探测信号频谱处理
# 包括原始信号的处理和补零后的处理（为了提高频率分辨力）

from FFT_Interpolation import line_cal, line_cal_fix, FFT_cal
import sys
import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
from matplotlib.ticker import (MultipleLocator, FormatStrFormatter,
                               AutoMinorLocator)
import matplotlib
from plt_style import *

colors = color_set_medium


def dual_Gaussian(x, sig, n_mid=0, sigma=0.2, pix_num=1280):
    N_mid_1 = 640 - n_mid
    N_mid_2 = 640 + n_mid
    gaussian_1 = np.exp(-0.5 * ((x - N_mid_1) / (sigma * pix_num / 2)) ** 2)
    gaussian_2 = np.exp(-0.5 * ((x - N_mid_2) / (sigma * pix_num / 2)) ** 2)
    window = (gaussian_1 + gaussian_2)
    #     c = np.exp(-0.5 * ((x-640)/((1/np.sqrt(2*n_mid))*1280))**2)
    c = np.exp(-0.5 * ((x - pix_num/2) / (sigma * pix_num / 2)) ** 2 * (12 * n_mid / pix_num)) * (1 - 2 * n_mid / pix_num) ** 2
    # print(c)
    # 信号
    data = ((sig - np.mean(sig)) * c + np.mean(sig)) * window
    return data


j = complex(0, 1)
c = 3e8  # 光速 [m/s]
Lamda = 633e-9  # 光波长 [m]
Fc = c / Lamda  # 光频率 [Hz]

pix_size = 5.3e-6  # d_p 像素尺寸
pix_num = 1280  # N_p 一行/列像素点总数
screen_diameter = pix_num * pix_size

x = np.linspace(0, (pix_num - 1), pix_num)

tau0 = pix_size
# 采样频率1/像素尺寸（类似空间频率）
fs = 1 / tau0
N = len(x)
# 信号频率3kHz
f = 3e3
# 相位
phi = 0
# 窗函数的中心偏移量
d_gauss = 0
# 相机量化的信号
sig = (1000 * np.cos(2 * np.pi * f * x * pix_size + phi) / 2 + 512) / 2
# 窗函数的中心值
N_mid_1 = 640 - d_gauss / 2
N_mid_2 = 640 + d_gauss / 2
# sigma的取值怎么定？
sigma = 0.2
# 高斯窗函数
gaussian_1 = np.exp(-0.5 * ((x - N_mid_1) / (sigma * pix_num / 2)) ** 2)
gaussian_2 = np.exp(-0.5 * ((x - N_mid_2) / (sigma * pix_num / 2)) ** 2)
window = (gaussian_1 + gaussian_2)

# plt.plot(window)
# 窗函数下的信号
data = dual_Gaussian(x=x, sig=sig, n_mid=d_gauss, sigma=sigma, pix_num=pix_num)
# 补零操作
zero_num = pix_num * 99
zero_padd = np.zeros(int(zero_num))
# 数据拼接
data_padd = np.concatenate((data, zero_padd))
# Fourier变换,包括补零与未补零
freqline, sig_FFT, sig_magnitude, sig_phase = FFT_cal(data, tau0)
freqline_padd, sig_FFT_padd, sig_magnitude_padd, sig_phase_padd = FFT_cal(data_padd, tau0)
# 观测两种情况下的频率分辨力
print(freqline[2] - freqline[1])
print(freqline_padd[2] - freqline_padd[1])

# plt.plot(freqline, sig_magnitude)
# plt.plot(freqline_padd, sig_magnitude_padd, '--r')

fig = plt.figure('Fig_3_3')
# 屏幕显示分辨率设定15:7
plt.gcf().set_size_inches(15 / 2.54, 7 / 2.54)
ax1 = fig.add_subplot(1, 2, 1)
ax2 = fig.add_subplot(1, 2, 2)
# freqline_padd  /1000是为了把空间频率单位由1/m→1/mm
# sig_magnitude_padd  ×100是为了补偿补零后傅里叶变换幅值的衰减
ax1.plot(freqline_padd / 1000, sig_magnitude_padd * 100, linewidth=1, color=color_set_normal['blue'], label='补零后;')
# 中轴线标注
ax1.plot([3, 3], [-20, 200], color=colors['red'], linestyle='dashed')
(markerLines, stemLines, baseLines) = ax1.stem(freqline / 1000, sig_magnitude, label='补零前', use_line_collection=True)
plt.setp(markerLines, marker='o', markersize=4, markeredgewidth=0, markerfacecolor=colors['black'])
plt.setp(baseLines, linewidth=0.5, color=colors['black'])
plt.setp(stemLines, linewidth=0.5, color=colors['black'])

ax1.set_xlim(0, 5)
ax1.set_xlabel('空间频率 $\mathrm{(/mm)}$')
ax1.set_ylabel('幅度 $\mathrm{(a.u.)}$')
ax1.legend(loc='upper right', ncol=2, handletextpad=0.3, columnspacing=0.5)

ax2.plot(freqline_padd / 1000, sig_phase_padd, color=color_set_normal['blue'], linewidth=1, label='补零后;')
(markerLines, stemLines, baseLines) = ax2.stem(freqline / 1000, sig_phase, label='补零前', use_line_collection=True)
plt.setp(markerLines, marker='o', markersize=4, markeredgewidth=0, markerfacecolor=colors['black'])
plt.setp(baseLines, linewidth=0.5, color=colors['black'])
plt.setp(stemLines, linewidth=0.5, color=colors['black'])
ax2.plot([3, 3], [-np.pi, np.pi], color=colors['red'], linestyle='dashed')
ax2.set_xlim(2.5, 3.5)
ax2.set_ylim(-3.5, 5)
ax2.set_xlabel('空间频率 $\mathrm{(/mm)}$')
ax2.set_ylabel('相位 $\mathrm{(rad)}$')
ax2.legend(loc='upper right', ncol=2, handletextpad=0.3, columnspacing=0.5)

ax1.text(-0.05, -0.1, '$\mathrm{a)}$', horizontalalignment='right', verticalalignment='top', transform=ax1.transAxes)
ax2.text(-0.05, -0.1, '$\mathrm{b)}$', horizontalalignment='right', verticalalignment='top', transform=ax2.transAxes)

for ax in [ax1, ax2]:
    ax.grid(axis='y', linestyle='-', alpha=0.3)
    labels = ax.get_xticklabels() + ax.get_yticklabels()
    [label.set_fontname('Times New Roman') for label in labels]
    ax.tick_params(axis='both', length=2, width=1, color=colors['grey'])

plt.tight_layout()
plt.show()
# plt.savefig(r'C:\Users\yu03\Dropbox\Ph_D\Figures\Fig_3_3.jpg', dpi=600) #指定分辨率保存
# plt.savefig(r'/Users/yl/Dropbox/Ph_D/Figures/Fig_3_3.jpg', dpi=600) #指定分辨率保存
